import os
import sys
import argparse

# 导入必要的函数
parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
sys.path.insert(0, parent_path)
from models.mmdetection._base_.base import infer
from dataset_utils import load_dataset_cfgs

# 定义推理的主函数
def main():
    parser = argparse.ArgumentParser(description='Inference by trained mask-rcnn detector')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('checkpoint', help='trained model\'s checkpoint')
    parser.add_argument('input_dir', help='input images directory')
    parser.add_argument('output_dir', help='output inference results')
    args = parser.parse_args()

    # 定义模型的推理配置文件
    base_cfg_path = os.path.join( parent_path, 'configs/yolo/yolov3_d53_8xb8-ms-608-273e_coco.py' )
    num_classes, user_cfg_dict = load_dataset_cfgs( args.config )

    # 我们还需要更改 head 中的 num_classes 以匹配数据集中的类别数
    user_cfg_dict['model'] = dict( bbox_head=dict(num_classes=num_classes) )

    # 调用推理函数
    infer( base_cfg_path, user_cfg_dict, args.checkpoint, args.input_dir, args.output_dir )

main()

